import cv2
import time
import torch
import socket
import numpy as np

# ESP32_CAM_URL = "http://192.168.137.100:81/stream"
# YOLO_MODEL_PATH = "runs/train/exp4/weights/best.pt"

# **ESP32-CAM 视频流 URL**
ESP32_CAM_URL = "http://192.168.43.15:81/stream"  # 请确保 IP 地址正确

# **YOLOv5 模型路径**
YOLO_MODEL_PATH = "E:\\yolov5-6.0\\best.pt"  # 你的 YOLOv5 训练模型

# 无人机
pyDrone_IP = "192.168.137.199"  # ESP32 的固定 IP
pyDrone_PORT = 1234             # ESP32 开放的端口

# 建立 UDP 连接
try:
    udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    udp_socket.connect((pyDrone_IP, pyDrone_PORT))
    print("成功连接 pyDrone")
except Exception as e:
    print("无法连接 pyDrone：", e)
    udp_socket = None

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

model = torch.hub.load('.', 'custom', path=YOLO_MODEL_PATH, source='local', device=device)
model.eval()


def iou(box1, box2):
    """计算两个框的 IoU（交并比）"""
    x1, y1, x2, y2 = box1
    x1_, y1_, x2_, y2_ = box2
    xi1, yi1 = max(x1, x1_), max(y1, y1_)
    xi2, yi2 = min(x2, x2_), min(y2, y2_)
    inter_area = max(0, xi2 - xi1) * max(0, yi2 - yi1)
    box1_area = (x2 - x1) * (y2 - y1)
    box2_area = (x2_ - x1_) * (y2_ - y1_)
    union_area = box1_area + box2_area - inter_area
    return inter_area / union_area if union_area > 0 else 0


def draw_detection(frame, boxes):
    """封装绘制边界框和标签的函数，返回目标中心坐标"""
    target_centers = []  # 存储每个目标的中心坐标
    for box in boxes:
        if isinstance(box, torch.Tensor):
            x1, y1, x2, y2, conf, cls = box.tolist()
        else:
            x1, y1, x2, y2, conf, cls = box
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)

        # 计算目标中心坐标
        center_x = (x1 + x2) / 2
        center_y = (y1 + y2) / 2
        target_centers.append((center_x, center_y))

        # 绘制边界框和标签
        label = f"Class {int(cls)}: {conf:.2f}"
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        # 绘制中心点
        cv2.circle(frame, (int(center_x), int(center_y)), 5, (0, 0, 255), -1)  # 红色圆点表示目标中心

    return frame, target_centers


def main():
    cap = cv2.VideoCapture(ESP32_CAM_URL)
    if not cap.isOpened():
        print("无法连接到 ESP32-CAM")
        return

    frame_count = 0
    process_interval = 2
    last_time = time.time()
    tracked_boxes = []

    while True:
        ret, frame = cap.read()
        if not ret:
            print("读取失败，尝试重新连接...")
            cap.release()
            cap = cv2.VideoCapture(ESP32_CAM_URL)
            if not cap.isOpened():
                print("重连失败，退出程序")
                break
            time.sleep(1)
            continue

        # 获取画面尺寸和中心
        frame_height, frame_width = frame.shape[:2]
        frame_center_x = frame_width / 2
        frame_center_y = frame_height / 2
        # 绘制画面中心点
        cv2.circle(frame, (int(frame_center_x), int(frame_center_y)), 5, (255, 0, 0), -1)  # 蓝色圆点表示画面中心

        frame_count += 1

        if frame_count % process_interval == 0:
            results = model(frame)
            current_boxes = results.xyxy[0]
            new_tracked_boxes = []
            for curr_box in current_boxes:
                curr_coords = curr_box[:4].tolist()
                best_match = None
                best_iou = 0
                for prev_box in tracked_boxes:
                    prev_coords = prev_box[:4]
                    iou_score = iou(curr_coords, prev_coords)
                    if iou_score > best_iou and iou_score > 0.5:
                        best_iou = iou_score
                        best_match = curr_box.tolist()
                new_tracked_boxes.append(best_match if best_match is not None else curr_box.tolist())
            tracked_boxes = new_tracked_boxes
            frame, target_centers = draw_detection(frame, tracked_boxes)
        elif tracked_boxes:
            frame, target_centers = draw_detection(frame, tracked_boxes)
        else:
            target_centers = []  # 没有目标时置为空

        # 计算并显示与画面中心的差值
        if target_centers:
            for i, (center_x, center_y) in enumerate(target_centers):
                diff_x = center_x - frame_center_x
                diff_y = center_y - frame_center_y
                diff_text = f"Target {i}: dx={diff_x:.1f}, dy={diff_y:.1f}"
                cv2.putText(frame, diff_text, (10, 60 + i * 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
                print(diff_text)  # 同时打印到控制台
                # 发送数据给 ESP32
                if udp_socket:
                    try:
                        msg = f"{diff_x:.1f},{diff_y:.1f}\n"
                        udp_socket.sendall(msg.encode())
                    except Exception as e:
                        print("⚠️ 发送失败:", e)

        # 计算并显示 FPS
        current_time = time.time()
        time_diff = current_time - last_time
        fps = 1 / time_diff if time_diff > 0 else 0
        last_time = current_time
        cv2.putText(frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 2)

        cv2.imshow("ESP32-CAM YOLOv5 Detection", frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    if udp_socket:
        udp_socket.close()


if __name__ == "__main__":
    main()